package tunnel

import (
	"encoding/json"
	"errors"
	"fmt"
	"golang.org/x/crypto/ssh"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
)

var (
	IdentUnique = make(map[string]bool, 4)
	ErrCh       = make(chan map[string]string, 4)
	ExCh        = make(chan map[string]bool, 4)
)

type Config struct {
	SSH     *SSHConfig        `json:"ssh"`
	Cfg     *ssh.ClientConfig //可以直接连ssh的一个实例化后的结构体
	Tunnels []*TConfig        `json:"tunnels"`
}

type TConfig struct {
	Ident      string            `json:"ident"`
	SSH        *SSHConfig        `json:"ssh"`
	Cfg        *ssh.ClientConfig //可以直接连ssh的一个实例化后的结构体
	LocalPort  int               `json:"local_port"`
	RemoteHost string            `json:"remote_host"`
	RemotePort int               `json:"remote_port"`
}

type SSHConfig struct {
	Host           string `json:"host"`
	User           string `json:"user"`
	Port           int    `json:"port"`
	Secret         string `json:"secret"`
	PrivateKeyFile string `json:"private_key"`
}

func (c *Config) Load(dir string) (err error) {
	var r io.ReadCloser
	var bytes []byte
	if r, err = os.Open(dir); err != nil {
		return
	}
	defer r.Close()
	bytes, err = ioutil.ReadAll(r)
	if err != nil {
		return
	}
	err = json.Unmarshal(bytes, c)
	if err != nil {
		return
	}
	c.Cfg = getClientConfig(c.SSH)
	for _, v := range c.Tunnels {
		if err = v.Valid(); err != nil {
			log.Printf("invalid Config, err=%v\n", err)
			return
		}
		if v.SSH == nil {
			v.SSH = c.SSH
			v.Cfg = c.Cfg
		} else {
			v.Cfg = getClientConfig(v.SSH)
		}
	}
	return
}

func (c *Config) Start(cfgDir *string) {
	err := c.Load(*cfgDir)
	if err != nil {
		log.Fatalf("config load err:%v", err)
	}
	log.Printf("配置加载完成，正在启动服务:\t%s\n", *cfgDir)
	for _, v := range c.Tunnels {
		go func(item *TConfig) {
			localAddr := HostPort("127.0.0.1", item.LocalPort)
			localListener, err := net.Listen("tcp", localAddr)
			if err != nil {
				log.Fatalf("[%-10s]net.Listen failed: %v", item.Ident, err)
			}
			log.Printf("[%-10s]%s\t1、listen succss\n", item.Ident, localAddr)
			for {
				localConn, err := localListener.Accept()
				if err != nil {
					if strings.Contains(err.Error(), "use of closed network connection") {
						log.Printf("[%-10s]2、%s\t: %v\n", item.Ident, "closed network", err)
						break
					}
					log.Fatalf("[%-10s]2、listen.Accept failed: %v\n", item.Ident, err)
				}
				log.Printf("[%-10s]%21s->%21s\t2、Connect succss\n", item.Ident, localConn.RemoteAddr().String(), localConn.LocalAddr().String())
				go item.Forward(localConn, item.Cfg)
			}
		}(v)
	}
	select {}
}

func getClientConfig(cfg *SSHConfig) *ssh.ClientConfig {
	var auth ssh.AuthMethod
	if cfg.PrivateKeyFile != "" {
		// true: privateKey specified
		auth = loadPrivateKeyFile(cfg.PrivateKeyFile)
	} else {
		auth = ssh.Password(cfg.Secret)
	}
	config := &ssh.ClientConfig{
		User: cfg.User,
		Auth: []ssh.AuthMethod{
			auth,
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}
	return config
}

func (t *TConfig) Valid() error {
	if t.Ident == "" {
		return errors.New("empty tunnel identify")
	}
	if _, ok := IdentUnique[t.Ident]; ok {
		return errors.New("duplicate tunnel ident=" + t.Ident)
	}
	IdentUnique[t.Ident] = true
	if t.LocalPort == 0 {
		// pass
	}
	if t.RemoteHost == "" {
		return errors.New("empty remote host")
	}
	if t.RemotePort == 0 {
		return errors.New("empty remote port")
	}
	return nil
}

func (t *TConfig) Forward(localConn net.Conn, config *ssh.ClientConfig) {
	defer func() {
		log.Printf("[%-10s]%21s->%21s\t7、local close\n", t.Ident, localConn.LocalAddr(), localConn.RemoteAddr())
		localConn.Close()
	}()
	//连接隧道服务器SSH
	tunnelAddr := HostPort(t.SSH.Host, t.SSH.Port)
	sshClientConn, err := ssh.Dial("tcp", tunnelAddr, config)
	if err != nil {
		log.Fatalf("[%s]3、ssh.Dial failed: %s", t.Ident, err)
	}
	localAddrSSH := localPort(sshClientConn.LocalAddr().String())
	log.Printf("[%-10s]%21s->%21s\t3、SSH.Dial succss\n", t.Ident, localAddrSSH, sshClientConn.RemoteAddr().String())
	defer func() {
		log.Printf("[%-10s]%21s->%21s\t6、ssh.connect close\n", t.Ident, localAddrSSH, sshClientConn.RemoteAddr().String())
		sshClientConn.Close()
	}()
	//隧道连接目标服务器
	targetAddr := HostPort(t.RemoteHost, t.RemotePort)
	sshConn, err := sshClientConn.Dial("tcp", targetAddr)
	if err != nil {
		log.Fatalf("[%s]4、ssh.Dial failed: %s", t.Ident, err)
	}
	log.Printf("[%-10s]%21s->%21s\t4、tunnel.Dial succss\n", t.Ident, tunnelAddr, targetAddr)
	defer func() {
		log.Printf("[%-10s]%21s->%21s\t5、tunnel.connect close\n", t.Ident, tunnelAddr, targetAddr)
		sshConn.Close()
	}()
	//本机<-->目标服务器 数据传送
	exist := make(chan bool)
	copyConn := func(dst, src net.Conn, ex chan bool) { //超时处理
		_, err = io.Copy(dst, src)
		if err != nil {
			msg := fmt.Sprintf("[%-10s]io.Copy failed: %v", t.Ident, err)
			log.Println(msg)
		}
		ex <- true
	}
	go copyConn(localConn, sshConn, exist)
	go copyConn(sshConn, localConn, exist)
	<-exist
}

func HostPort(host string, port int) string {
	return net.JoinHostPort(host, strconv.Itoa(port))
}

func loadPrivateKeyFile(dir string) ssh.AuthMethod {
	buffer, err := ioutil.ReadFile(dir)
	if err != nil {
		return nil
	}
	key, err := ssh.ParsePrivateKey(buffer)
	if err != nil {
		return nil
	}
	return ssh.PublicKeys(key)
}

func (t *TConfig) sendErr(msg string, err error) {
	ExCh <- map[string]bool{t.Ident: true}
	msg = fmt.Sprintf("%s:%v", msg, err)
	ErrCh <- map[string]string{t.Ident: msg}
	//log.Println(msg)
}

func localPort(addr string) string {
	_, port, _ := net.SplitHostPort(addr)
	return net.JoinHostPort("127.0.0.1", port)
}
